#include "sampler.hpp"
#include "../math/random_generator.hpp"
#include "../math/tvector.hpp"
#include <cmath>

#pragma warning( push )
#pragma warning( disable : 4244 )

/** 使用cubic
 * \note
 */
inline f32 solve( f32 r ) {
  f32 u = r;
  for ( ui32 i = 0; i < 5; i++ ) {
    u = ( 11.f * r + u * u * ( 6.f + u * ( 8.f - 9.f * u ) ) ) /
        ( 4.f + 12.f * u * ( 1.f + u * ( 1.f - u ) ) );
  }
  return u;
}

/** 三次多项式
 * \note 给定一个坐标，求解这个三次多项式
 * \param x the coordinate
 * \return a cubic interpolate value
 */
inline f32 cubicFilter( f32 x ) {
  // a formulator
  const f32 a = 1.f / 24.f;
  const f32 b = 23.f / 24.f;
  if ( x < a ) {
    return powf( 24.f * x, 0.25f ) - 2.f;
  } else if ( x < .5f ) {
    return solve( 24.f * ( x - a ) / 11.f ) - 1.f;
  } else if ( x < b ) {
    return 1.f - solve( 24.f * ( b - x ) / 11.f );
  } else
    return 2 - powf( 24.f * ( 1.f - x ), 0.25f );
}

void SampleAndFilteringFunc::regularGridSample( vec2f *samples,
                                                ui32 num_samples ) {
  if ( samples == nullptr ) return;
  ui32 sqrt_samples = ui32( sqrtf( (f32)num_samples ) );
  for ( ui32 i = 0; i < sqrt_samples; i++ )
    for ( ui32 j = 0; j < sqrt_samples; j++ ) {
      samples[i * sqrt_samples + j].x = ( (f32)i + 0.5f ) / (f32)sqrt_samples;
      samples[i * sqrt_samples + j].y = ( (f32)j + 0.5f ) / (f32)sqrt_samples;
    }
}

/**
 * 简单随机采样
 * \note 随机采样即为每个采样点生成一个随机数
 *
 * \param samples 原采样点
 * \param num_samples 原采样点的数量
 */
void SampleAndFilteringFunc::random( vec2f *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  mt_uniform_float_rg rg( -.5f, .5f );
  for ( ui32 i = 0; i < num_samples; i++ ) {
    samples[i].x = rg();
    samples[i].y = rg();
  }
}

/**
 * Jittered采样方法
 * \note 在一个像素中分成nxn的网格，在每个小格中随机生成一个采样点
 *
 * \param samples 原采样点
 * \param num_samples 原采样点的数量
 */
void SampleAndFilteringFunc::jittered( vec2f *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  mt_uniform_float_rg rg( -.5f, .5f );
  ui32 sqrt_samples = static_cast<ui32>( sqrtf( (f32)num_samples ) );
  for ( ui32 i = 0; i < sqrt_samples; i++ )
    for ( ui32 j = 0; j < sqrt_samples; j++ ) {
      // 生成随机点
      f32 x = ( ( f64 )( i ) + rg() ) / (f64)sqrt_samples;
      f32 y = ( ( f64 )( j ) + rg() ) / (f64)sqrt_samples;
      samples[i * sqrt_samples + j].x = x;
      samples[i * sqrt_samples + j].y = y;
    }
}

/**
 * \brief n-rooks采样 由Shirley在1991年提出
 * \note 采样步骤为
 * 1. 首先在单个采样单元中，沿着主对角线生成随机点
 * 2. 然后随机偏移，上一个生成的随机点，保证每个采样点的
 * 所在行列只有该采样点
 *
 * \param samples 原采样点
 * \param num_samples 原采样点的数量
 */
void SampleAndFilteringFunc::nrooks( vec2f *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  mt_uniform_float_rg rg( -.5f, .5f );
  for ( ui32 i = 0; i < num_samples; i++ ) {
    samples[i].x = ( (f64)i + rg() ) / (f64)num_samples;
    samples[i].y = ( (f64)i + rg() ) / (f64)num_samples;
  }
  // shuffle the x coordinate
  // 随机索引，并进行位置交换，以此来偏移x坐标
  for ( ui32 i = num_samples - 2; i >= 0; i-- ) {
    ui32 target = ui32( rg() * (f64)i );
    f32 temp = samples[i + 1].x;
    samples[i + 1].x = samples[target].x;
    samples[target].x = temp;
  }
}

/**
 * multi-jittered 采样，结合n-rooks以及jittered采样的优点由Chui et
 * al.在1994年提出
 * \note 采样步骤为：
 * 1. 先将该像素依据采样的点数量分成两级网格：一个是√nx√n的网格，
 *    另一个是nxn的子网格
 * 2. 然后依据n-rooks的条件为每个√nx√n网格中的每个格子生成一个采样点
 * 3. 进行xy坐标的随机偏移
 *
 * \param samples 原采样点
 * \param num_samples 原采样点的数量
 */
void SampleAndFilteringFunc::multiJittered( vec2f *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  mt_uniform_float_rg rg( -.5f, .5f );
  ui32 sqrt_sample = static_cast<ui32>( sqrt<ui32>( num_samples ) );
  f32 subcell_width = 1.f / (f32)num_samples;
  for ( ui32 i = 0; i < sqrt_sample; i++ )
    for ( ui32 j = 0; j < sqrt_sample; j++ ) {
      samples[i * sqrt_sample + j].x = i * sqrt_sample * subcell_width +
                                       j * subcell_width + rg() * subcell_width;
      samples[i * sqrt_sample + j].y = j * sqrt_sample * subcell_width +
                                       i * subcell_width + rg() * subcell_width;
    }

  // 进行xy坐标的随机偏移
  for ( ui32 i = 0; i < sqrt_sample; i++ )
    for ( ui32 j = 0; j < sqrt_sample; j++ ) {
      ui32 k = j + ui32( rg() * ( sqrt_sample - j - 1 ) );
      f32 t = samples[i * sqrt_sample + j].x;
      samples[i * sqrt_sample + j].x = samples[i * sqrt_sample + k].x;
      samples[i * sqrt_sample + k].x = t;

      k = j + ui32( rg() * ( sqrt_sample - j - 1 ) );
      t = samples[j * sqrt_sample + i].y;
      samples[j * sqrt_sample + i].y = samples[k * sqrt_sample + i].y;
      samples[k * sqrt_sample + i].y = t;
    }
}

void SampleAndFilteringFunc::shuffle( vec2f *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  mt_uniform_float_rg rg( -.5f, .5f );
  for ( ui32 i = num_samples - 2; i >= 0; i-- ) {
    ui32 target = ui32( rg() * (f64)i );
    auto temp = samples[i + 1];
    samples[i + 1] = samples[target];
    samples[target] = temp;
  }
}

/**
 * BoxFilter 该滤波器能够进行简单的图像重建
 * \note 采样步骤
 * 1. 使用每个像素点的中心作为新的采样点
 *
 * \param samples 原采样点
 * \param num_samples 原采样点的数量
 */
void SampleAndFilteringFunc::boxFilter( vec2f *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  for ( ui32 i = 0; i < num_samples; i++ ) {
    samples[i] = samples[i] - .5f;
  }
}

/**
 * TentFilter boxFilter的改进版本
 * \note 滤波步骤：
 * 1. 根据Linear b-spline 函数给定一个权重值
 * 2. 距离像素中心越近的权重越大
 */
void SampleAndFilteringFunc::tentFilter( vec2f *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  for ( ui32 i = 0; i < num_samples; i++ ) {
    f32 x = samples[i].x;
    f32 y = samples[i].y;

    if ( x < .5f )
      samples[i].x = (f32)sqrt( 2.0 * (f64)x ) - 1.f;
    else
      samples[i].x = ( 1.f - (f32)sqrt( 2.0 - 2.0 * (f64)x ) );

    if ( y < .5f )
      samples[i].y = (f32)sqrt( 2.0 * (f64)y ) - 1.f;
    else
      samples[i].y = ( 1.f - (f32)sqrt( 2.0 - 2.0 * (f64)y ) );
  }
}

/**
 * 三次spline滤波
 * \note cubic polynomial
 * 1. 通过一个三次spline函数采样信号
 *
 *
 */
void SampleAndFilteringFunc::cubicSplineFiter( vec2f *samples,
                                               ui32 num_samples ) {
  if ( samples == nullptr ) return;
  for ( ui32 i = 0; i < num_samples; i++ ) {
    samples[i] =
        vec2f( cubicFilter( samples[i].x ), cubicFilter( samples[i].y ) );
  }
}

void SampleAndFilteringFunc::random_1D( f32 *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  mt_uniform_float_rg rg( 0, 1 );
  for ( ui32 i = 0; i < num_samples; i++ ) samples[i] = rg();
}

void SampleAndFilteringFunc::jittered_1D( f32 *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  mt_uniform_float_rg rg( 0, 1 );
  for ( ui32 i = 0; i < num_samples; i++ )
    samples[i] = ( i + rg() ) / (f64)num_samples;
}

void SampleAndFilteringFunc::shuffle_1D( f32 *samples, ui32 num_samples ) {
  if ( samples == nullptr ) return;
  mt_uniform_float_rg rg( 0, 1 );
  for ( ui32 i = num_samples - 2; i >= 0; i-- ) {
    ui32 k = ui32( rg() * i );
    f32 t = samples[k];
    samples[k] = samples[i + 1];
    samples[i + 1] = t;
  }
}

#pragma warning( pop )